import time
import os
import matplotlib.pyplot as plt
import numpy as np

# 导入你实现的算法
from publickey.rsa_core import RSA
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import (
    get_curve_by_name,
    generate_keypair as ecc_gen_keys,
    ecc_encrypt_simplified,
    ecc_decrypt_simplified
)

def plot_results(results):
    """将实验结果可视化"""
    # 指定一个支持中文的字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei'] 
    plt.rcParams['axes.unicode_minus'] = False 

    # 修改: 定义新的颜色方案
    color_palettes = {
        'RSA': ['#ff9999', '#ff4d4d', '#cc0000'],  # 浅红, 红, 深红
        'ElGamal': ['#99ff99', '#4dff4d', '#009933'], # 浅绿, 绿, 深绿
        'ECC': ['#9999ff', '#4d4dff', '#0000cc']   # 浅蓝, 蓝, 深蓝
    }
    
    # --- 1. 密钥生成时间 ---
    labels = list(results['key_gen_time'].keys())
    values = list(results['key_gen_time'].values())
    
    # 修改: 动态生成颜色列表
    colors_key_gen = []
    counters = {'RSA': 0, 'ElGamal': 0, 'ECC': 0}
    for label in labels:
        algo_family = label.split('-')[0]
        colors_key_gen.append(color_palettes[algo_family][counters[algo_family]])
        counters[algo_family] += 1
        
    plt.figure(figsize=(12, 8))
    bars = plt.bar(labels, values, color=colors_key_gen)
    
    # 修改: Y轴标签单位为毫秒
    plt.ylabel('时间 (毫秒)')
    plt.title('密钥生成时间比较')
    plt.yscale('log')
    plt.xticks(rotation=45, ha="right")
    
    # 修改: 在柱形图上显示数值
    plt.bar_label(bars, fmt='%.2f', padding=3, fontsize=9)
    
    plt.tight_layout()
    plt.savefig('key_generation_time.png')
    plt.show()

    # --- 2. 加密/解密时间 (平均值) ---
    labels = list(results['encrypt_time'].keys())
    encrypt_times = list(results['encrypt_time'].values())
    decrypt_times = list(results['decrypt_time'].values())
    
    x = np.arange(len(labels))
    width = 0.35
    
    fig, ax = plt.subplots(figsize=(14, 9))
    
    # 修改: 使用与上面相同的动态颜色列表
    rects1 = ax.bar(x - width/2, encrypt_times, width, label='加密', color='skyblue')
    rects2 = ax.bar(x + width/2, decrypt_times, width, label='解密', color='salmon')
    
    # 修改: Y轴标签单位为毫秒
    ax.set_ylabel('平均时间 (毫秒)')
    ax.set_title('平均加密与解密时间比较 (100次)')
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.legend()
    ax.set_yscale('log')

    # 修改: 在柱形图上显示数值
    ax.bar_label(rects1, fmt='%.4f', padding=3, fontsize=8, rotation=90)
    ax.bar_label(rects2, fmt='%.4f', padding=3, fontsize=8, rotation=90)
    
    plt.tight_layout()
    plt.savefig('encryption_decryption_time_avg.png')
    plt.show()

def run_experiments():
    """运行所有性能评估实验"""
    
    num_iterations = 100
    print(f"每个算法将进行 {num_iterations} 次加密/解密操作来计算平均时间。")
    
    # --- 参数设置 ---
    rsa_elgamal_key_sizes = [1024, 2048, 3072] 
    ecc_curve_names = ["secp256k1", "secp256r1", "secp384r1"]
    message_size_bytes = 40

    results = {
        "key_gen_time": {},
        "encrypt_time": {},
        "decrypt_time": {},
        "ciphertext_size": {}
    }

    # --- 1. RSA 实验 ---
    print("\n" + "="*10, "开始 RSA 实验", "="*10)
    messages_rsa = [os.urandom(message_size_bytes) for _ in range(num_iterations)]
    for bits in rsa_elgamal_key_sizes:
        print(f"\n--- 测试 RSA with {bits}-bit key ---")
        algo_id = f"RSA-{bits}"
        start_time = time.perf_counter()
        rsa_instance = RSA(bits=bits)
        # 修改: 结果乘以1000，单位变为毫秒
        results["key_gen_time"][algo_id] = (time.perf_counter() - start_time) * 1000
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.4f} 毫秒")

        start_time = time.perf_counter()
        ciphertexts_rsa = [rsa_instance.encrypt(m) for m in messages_rsa]
        results["encrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 毫秒")

        start_time = time.perf_counter()
        for c in ciphertexts_rsa: rsa_instance.decrypt(c)
        results["decrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 毫秒")
        
        n, _ = rsa_instance.public_key
        results["ciphertext_size"][algo_id] = (message_size_bytes, (n.bit_length() + 7) // 8)

    # --- 2. ElGamal 实验 ---
    print("\n" + "="*10, "开始 ElGamal 实验", "="*10)
    messages_elgamal_int = [int.from_bytes(os.urandom(message_size_bytes), 'big') for _ in range(num_iterations)]
    for bits in rsa_elgamal_key_sizes:
        print(f"\n--- 测试 ElGamal with {bits}-bit key ---")
        algo_id = f"ElGamal-{bits}"
        if bits > 2048:
            print("警告: 3072-bit ElGamal密钥生成可能非常耗时，请耐心等待...")
        start_time = time.perf_counter()
        elgamal_instance = ElGamal(key_size_bits=bits)
        results["key_gen_time"][algo_id] = (time.perf_counter() - start_time) * 1000
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.4f} 毫秒")
        
        pub_key, priv_key = elgamal_instance.public_key, elgamal_instance.private_key
        public_params = (pub_key[0], pub_key[1], pub_key[2])
        
        start_time = time.perf_counter()
        ciphertexts_elgamal = [ElGamal.encrypt(pub_key, m) for m in messages_elgamal_int]
        results["encrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 毫秒")
        
        start_time = time.perf_counter()
        for c in ciphertexts_elgamal: ElGamal.decrypt(priv_key, public_params, c)
        results["decrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 毫秒")
        
        p, _, _, _ = pub_key
        results["ciphertext_size"][algo_id] = (message_size_bytes, 2 * ((p.bit_length() + 7) // 8))

    # --- 3. ECC 实验 ---
    print("\n" + "="*10, "开始 ECC 实验", "="*10)
    messages_ecc = [os.urandom(message_size_bytes) for _ in range(num_iterations)]
    for curve_name in ecc_curve_names:
        print(f"\n--- 测试 ECC with {curve_name} ---")
        algo_id = f"ECC-{curve_name}"

        try:
            curve = get_curve_by_name(curve_name)
        except ValueError as e:
            print(e); continue

        start_time = time.perf_counter()
        ecc_priv_key, ecc_pub_key = ecc_gen_keys(curve)
        results["key_gen_time"][algo_id] = (time.perf_counter() - start_time) * 1000
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.4f} 毫秒")

        start_time = time.perf_counter()
        ciphertexts_ecc = [ecc_encrypt_simplified(ecc_pub_key, curve, m) for m in messages_ecc]
        results["encrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 毫秒")

        start_time = time.perf_counter()
        for R, C in ciphertexts_ecc: ecc_decrypt_simplified(ecc_priv_key, curve, R, C)
        results["decrypt_time"][algo_id] = ((time.perf_counter() - start_time) / num_iterations) * 1000
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 毫秒")

        point_size_bytes = 1 + 2 * ((curve.p.bit_length() + 7) // 8)
        cipher_len = point_size_bytes + message_size_bytes 
        results["ciphertext_size"][algo_id] = (message_size_bytes, cipher_len)
    
    return results

if __name__ == "__main__":
    experiment_results = run_experiments()
    # 打印最终结果摘要
    print("\n\n" + "="*20, "实验结果摘要", "="*20)
    print("\n密钥生成时间 (毫秒):")
    for k, v in experiment_results["key_gen_time"].items():
        print(f"  {k:<15}: {v:,.2f} ms")
    
    print("\n平均加密时间 (毫秒):")
    for k, v in experiment_results["encrypt_time"].items():
        print(f"  {k:<15}: {v:.4f} ms")

    print("\n平均解密时间 (毫秒):")
    for k, v in experiment_results["decrypt_time"].items():
        print(f"  {k:<15}: {v:.4f} ms")
        
    print("\n数据扩展性 (明文 -> 密文):")
    for k, v in experiment_results["ciphertext_size"].items():
        plain_size, cipher_size = v
        # 修改: 计算并添加扩展百分比
        expansion_ratio = cipher_size / plain_size
        print(f"  {k:<15}: {plain_size} -> {cipher_size} 字节 (扩展为 {expansion_ratio:.1%})")
        
    plot_results(experiment_results)